package com.lovely602.ai.alibaba.controller.documentretriever;


import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeDocumentRetriever;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeDocumentRetrieverOptions;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

/**
 * 文档检索
 *
 * @author lizhixing
 */
@Slf4j
@RestController
@RequestMapping(value = "/document")
public class DocumentRetrieverController {

    private final DocumentRetriever documentRetriever;

    public DocumentRetrieverController() {
        DashScopeApi dashScopeApi = DashScopeApi.builder()
                .apiKey(System.getenv("AI_DASHSCOPE_API_KEY"))
                .build();

        DashScopeDocumentRetrieverOptions retrieverOptions = DashScopeDocumentRetrieverOptions.builder()
                .withIndexName("spring")
                .build();

        this.documentRetriever = new DashScopeDocumentRetriever(dashScopeApi, retrieverOptions);
    }

    /**
     * 文档检索
     */
    @GetMapping(value = "/retriever")
    public List<Document> retriever(@RequestParam(value = "msg", defaultValue = "What's spring ai?", required = true) String msg) {

        Query query = new Query(msg);

        return documentRetriever.retrieve(query);
    }


}
